from environment.ground_user import GroundUser
from environment.uav import UAV
import numpy as np
import math

energy_consume = 1


class MultiUAVEnv:
    def __init__(self, num_uavs, grid_size, num_gu,num_position_action,num_energy_action,charging_stations=[(500, 500),(200,200),(800,800)]):
        self.grid_size = grid_size
        self.ground_users = [GroundUser(grid_size) for _ in range(num_gu)]
        self.uavs = [UAV(grid_size,i,charging_stations,num_position_action,num_energy_action) for i in range(num_uavs)]
        self.charging_stations = charging_stations

    def get_global_state(self):
        return np.concatenate([uav.get_state() for uav in self.uavs])

    def compute_joint_reward(self):
        reward = 0
        for gu in self.ground_users:
            reward += 10 if gu.served else -1
        return reward
    
    def perform_action(self, uav, energy_action, cruise_action):
        reward = 0
        reward_energy = 0
        charging_accept = True
        
        charging_uavs =  [other_uav for other_uav in self.uavs if other_uav.position in self.charging_stations and other_uav != uav]
        
        if len(charging_uavs) == len(self.uavs) -1:
            charging_accept = False
        else:
            charging_accept = True
            
        # 检查无人机是否在充电站位置，且选择了充电动作
        if uav.position in self.charging_stations and energy_action and charging_accept== 0:
            uav.energy = 100  # 重置电量为满电状态
            print(f"UAV{uav.id} has already recharged and cannot perform other actions.")
            return -1,0# 返回0奖励和当前状态
        else:
            uav.reduce_energy(energy_consume)  # 每个非充电动作都消耗固定电量
            reward -= energy_consume
            
        if energy_action:
            print(f"UAV {uav.id} at {uav.position} choose the energy action 1. Heading to nearest charging station.")
            uav.go_to_nearest_charging_station ()
        else:
            # 执行飞行控制动作
            reward += uav.execute_cruise_action(cruise_action)
            reward_energy += 1
                
        service_providing_nows = []
        
        
        #检查是否有用户在当前位置，并提供服务
        service_reward = 0
        for gu in self.ground_users:
            if not gu.served and np.linalg.norm(np.array(gu.position)-np.array(uav.position))<= 100:
                print(f"UAV {uav.id} at {uav.position} is near user at {gu.position} with demand {gu.communication_demand}")
                service_reward = uav.serve_ground_user(gu.position,gu.communication_demand)
                if service_reward == 1:
                    # reward = reward + gu.communication_demand/2
                    gu.served = True
                    service_providing_nows.append(gu.communication_demand)
                    #print(f"UAV {uav.id} at {uav.position} successfully served user at {gu.position}")
                    reward -= 3
                else:
                    service_providing_nows.append(0)
            else:
                    service_providing_nows.append(0)
        
        if len(service_providing_nows) == 0:
            for gu in self.ground_users:
                service_providing_nows.append(0)
        
        for other_uav in self.uavs:
            if other_uav != uav:
                if np.linalg.norm(np.array(other_uav.position)-np.array(uav.position))<= 100:
                    print(f"UAV {uav.id} at {uav.position} is near UAV {other_uav.id} at {other_uav.position}")
                    reward -= 5
        
        
        fairness_reward = self.calculate_fairness_reward(service_providing_nows)
        other_reward = self.calculate_other_reward(service_providing_nows)
        reward += fairness_reward*other_reward/energy_consume
        #print(f"the fainess reward is {fairness_reward} and the other reward is {other_reward}")
        
        return reward_energy,reward
    def calculate_fairness_reward(self,service_providing_nows):
        # 计算服务公平性奖励
        sum_demand = 0
        b_k = []
        mu_k = []
        b_k_sum = 0
        beta_k_sum = 0
        beta_k = []
        fenzi = 0
        fenmu = 0
        i = 0
        for gu in self.ground_users:
            sum_demand += gu.communication_demand
            
        for gu in self.ground_users:
            b_k.append(gu.communication_demand/sum_demand)
            b_k_sum +=gu.communication_demand/sum_demand
        
        for gu,bk in zip(self.ground_users,b_k):
            mu_k1 = (1+1/(math.log(b_k[1]/b_k_sum)))/(1+1/(math.log(bk/b_k_sum)))
            mu_k.append(mu_k1)
        
        
        for gu_served in service_providing_nows:
            beta_k_sum += gu_served
            
            
        if beta_k_sum >0:
            print("kaishitiaoshi")
            
        for gu_served in service_providing_nows:
            if beta_k_sum == 0:
                beta_k.append(0)
            else:
                beta_k.append(gu_served/beta_k_sum)
            
        for gu in self.ground_users :
            if beta_k[i] == 0:
                fenzi +=0
            else:
                fenzi += mu_k[i]*beta_k[i]*(math.log(1+beta_k[i])/math.log(b_k[i]+1))*gu.communication_demand
            fenmu += mu_k[i]*gu.communication_demand
            i+=1
            

        fairness_reward = fenzi/fenmu
        return fairness_reward

    def calculate_other_reward(self,service_providing_nows):
        other_reaward = 0
        fenmu = 0
        for uav_served in service_providing_nows:
            if uav_served > 0:
                fenmu += 1
        if fenmu == 0:
            return 0
        for uav_served in service_providing_nows:
            fenzi = uav_served
            other_reaward += fenzi/fenmu
            
        return other_reaward

    
    
    